"""
LangGraph工作流管理器
定义和管理各种工作流
"""
from typing import Dict, Any, List, Optional
from langgraph.graph import StateGraph, END
from app.core.logging import get_logger
from app.graphs.core import GraphState, WorkflowConfig
from app.graphs.nodes import NODE_MAPPING

logger = get_logger(__name__)


class WorkflowManager:
    """工作流管理器"""
    
    def __init__(self):
        self.workflows = {}
        self._initialize_workflows()
    
    def _initialize_workflows(self):
        """初始化工作流"""
        # 简单聊天工作流
        self.workflows['simple_chat'] = self._create_simple_chat_workflow()
        
        # 工具增强工作流
        self.workflows['tool_enhanced'] = self._create_tool_enhanced_workflow()
        
        # 推理链工作流
        self.workflows['reasoning'] = self._create_reasoning_workflow()
        
        # 流式响应工作流
        self.workflows['streaming'] = self._create_streaming_workflow()
    
    def _create_simple_chat_workflow(self) -> StateGraph:
        """创建简单聊天工作流"""
        workflow = StateGraph(GraphState)
        
        # 添加节点
        workflow.add_node("validate_input", NODE_MAPPING["validate_input"])
        workflow.add_node("generate_response", NODE_MAPPING["generate_response"])
        workflow.add_node("error_handling", NODE_MAPPING["error_handling"])
        workflow.add_node("complete", NODE_MAPPING["complete"])
        
        # 设置入口点
        workflow.set_entry_point("validate_input")
        
        # 添加边（路由逻辑）
        workflow.add_edge("validate_input", "generate_response")
        
        # 从generate_response到complete或error_handling
        workflow.add_conditional_edges(
            "generate_response",
            self._route_from_generate_response,
            {
                "complete": "complete",
                "error_handling": "error_handling",
                "retry": "generate_response"
            }
        )
        
        # 从error_handling到complete或generate_response
        workflow.add_conditional_edges(
            "error_handling",
            self._route_from_error_handling,
            {
                "complete": "complete",
                "generate_response": "generate_response"
            }
        )
        
        workflow.add_edge("complete", END)
        
        return workflow.compile()
    
    def _create_tool_enhanced_workflow(self) -> StateGraph:
        """创建工具增强工作流"""
        workflow = StateGraph(GraphState)
        
        # 添加节点
        workflow.add_node("validate_input", NODE_MAPPING["validate_input"])
        workflow.add_node("tool_selection", NODE_MAPPING["tool_selection"])
        workflow.add_node("tool_execution", NODE_MAPPING["tool_execution"])
        workflow.add_node("generate_response", NODE_MAPPING["generate_response"])
        workflow.add_node("error_handling", NODE_MAPPING["error_handling"])
        workflow.add_node("complete", NODE_MAPPING["complete"])
        
        # 设置入口点
        workflow.set_entry_point("validate_input")
        
        # 从validate_input到tool_selection
        workflow.add_edge("validate_input", "tool_selection")
        
        # 从tool_selection到tool_execution或generate_response
        workflow.add_conditional_edges(
            "tool_selection",
            self._route_from_tool_selection,
            {
                "tool_execution": "tool_execution",
                "generate_response": "generate_response"
            }
        )
        
        # 从tool_execution到generate_response
        workflow.add_edge("tool_execution", "generate_response")
        
        # 从generate_response到complete或error_handling
        workflow.add_conditional_edges(
            "generate_response",
            self._route_from_generate_response,
            {
                "complete": "complete",
                "error_handling": "error_handling",
                "retry": "generate_response"
            }
        )
        
        # 从error_handling到complete或generate_response
        workflow.add_conditional_edges(
            "error_handling",
            self._route_from_error_handling,
            {
                "complete": "complete",
                "generate_response": "generate_response"
            }
        )
        
        workflow.add_edge("complete", END)
        
        return workflow.compile()
    
    def _create_reasoning_workflow(self) -> StateGraph:
        """创建推理链工作流"""
        workflow = StateGraph(GraphState)
        
        # 添加节点
        workflow.add_node("validate_input", NODE_MAPPING["validate_input"])
        workflow.add_node("reasoning_chain", NODE_MAPPING["generate_response"])  # 复用生成节点
        workflow.add_node("generate_response", NODE_MAPPING["generate_response"])
        workflow.add_node("error_handling", NODE_MAPPING["error_handling"])
        workflow.add_node("complete", NODE_MAPPING["complete"])
        
        # 设置入口点
        workflow.set_entry_point("validate_input")
        
        # 添加边
        workflow.add_edge("validate_input", "reasoning_chain")
        workflow.add_edge("reasoning_chain", "generate_response")
        
        # 条件路由
        workflow.add_conditional_edges(
            "generate_response",
            self._route_from_generate_response,
            {
                "complete": "complete",
                "error_handling": "error_handling"
            }
        )
        
        workflow.add_conditional_edges(
            "error_handling",
            self._route_from_error_handling,
            {
                "complete": "complete",
                "generate_response": "generate_response"
            }
        )
        
        workflow.add_edge("complete", END)
        
        return workflow.compile()
    
    def _create_streaming_workflow(self) -> StateGraph:
        """创建流式响应工作流"""
        workflow = StateGraph(GraphState)
        
        # 添加节点
        workflow.add_node("validate_input", NODE_MAPPING["validate_input"])
        workflow.add_node("generate_response", NODE_MAPPING["generate_response"])
        workflow.add_node("stream_processing", NODE_MAPPING["generate_response"])  # 复用生成节点
        workflow.add_node("complete", NODE_MAPPING["complete"])
        
        # 设置入口点
        workflow.set_entry_point("validate_input")
        
        # 添加边
        workflow.add_conditional_edges(
            "validate_input",
            self._route_for_streaming,
            {
                "streaming": "stream_processing",
                "normal": "generate_response"
            }
        )
        
        workflow.add_edge("stream_processing", "complete")
        workflow.add_edge("generate_response", "complete")
        workflow.add_edge("complete", END)
        
        return workflow.compile()
    
    # 路由函数
    def _route_from_generate_response(self, state: GraphState) -> str:
        """从生成响应节点路由"""
        if state.get("error_occurred"):
            if state.get("retry_count", 0) < 3:
                return "retry"
            else:
                return "error_handling"
        else:
            return "complete"
    
    def _route_from_error_handling(self, state: GraphState) -> str:
        """从错误处理节点路由"""
        if state.get("error_occurred"):
            return "complete"  # 结束工作流
        else:
            return "generate_response"  # 重试
    
    def _route_from_tool_selection(self, state: GraphState) -> str:
        """从工具选择节点路由"""
        if state.get("tools_used"):
            return "tool_execution"
        else:
            return "generate_response"
    
    def _route_for_streaming(self, state: GraphState) -> str:
        """流式响应路由"""
        if state.get("is_streaming"):
            return "streaming"
        else:
            return "normal"
    
    def get_workflow(self, workflow_type: str) -> Optional[StateGraph]:
        """获取工作流"""
        return self.workflows.get(workflow_type)
    
    def list_available_workflows(self) -> List[str]:
        """获取可用工作流列表"""
        return list(self.workflows.keys())
    
    async def execute_workflow(
        self,
        workflow_type: str,
        initial_state: Dict[str, Any]
    ) -> Dict[str, Any]:
        """执行工作流"""
        workflow = self.get_workflow(workflow_type)
        
        if not workflow:
            raise ValueError(f"不支持的工流类型: {workflow_type}")
        
        try:
            # 执行工作流
            logger.info(f"开始执行工作流: {workflow_type}")
            
            final_state = await workflow.ainvoke(initial_state)
            
            # 提取结果
            result = {
                "workflow_type": workflow_type,
                "success": not final_state.get("error_occurred", False),
                "messages": final_state.get("messages", []),
                "total_tokens": final_state.get("total_tokens", 0),
                "estimated_cost": final_state.get("estimated_cost", 0.0),
                "tools_used": final_state.get("tools_used", []),
                "function_calls": final_state.get("function_calls", []),
                "execution_time": None
            }
            
            # 计算执行时间
            if final_state.get("start_time") and final_state.get("end_time"):
                result["execution_time"] = (
                    final_state["end_time"] - final_state["start_time"]
                ).total_seconds()
            
            if final_state.get("error_occurred"):
                result["error_message"] = final_state.get("error_message")
            
            logger.info(f"工作流执行完成: {workflow_type}")
            
            return result
        
        except Exception as e:
            logger.error(f"工作流执行失败: {e}")
            raise


# 全局工作流管理器实例
workflow_manager = WorkflowManager()